'''
Created on May 20, 2009

@author: mkiyer
'''

from sqlalchemy import Table, Column, Integer, String, Text, MetaData, \
    ForeignKey, and_, or_, not_, create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import aliased, relation, backref, sessionmaker
from sqlalchemy.sql import func

from veggie.genome.interval import is_overlapping

Base = declarative_base()

class Annotation(Base):
    '''
    represents a genome annotation
    
    Fields ---
    chrom: chromosome in UCSC format e.g. 'chr1', 'chrX', etc
    chromStart: starting location on chromosome, numbering starts from zero
    chromEnd: ending location on chromosome, exclusive e.g. the region (0, 100) includes bases 0 through 99
    name: name of this annotation e.g. TMPRSS2, ChIPseq_ERG_HPeak, etc
    type: type of annotation (gene, protein binding site, motif, etc)
    '''
    __tablename__ = 'annotations'

    id = Column(Integer, primary_key=True)
    chrom = Column(String(255))
    chromStart = Column(Integer)
    chromEnd = Column(Integer)
    name = Column(Text, nullable=False)
    type = Column(Text)
    
    def __init__(self, chrom, chrom_start, chrom_end, name, type):
        self.chrom = chrom
        self.chromStart = chrom_start
        self.chromEnd = chrom_end
        self.name = name
        self.type = type

    def __cmp__(self, other):
        if self.__eq__(other):
            return 0
        else:
            return self.id - other.id

    def __eq__(self, other):
        return (self.chrom == other.chrom) and \
               (self.chromStart == other.chromStart) and \
               (self.chromEnd == other.chromEnd) and \
               (self.name == other.name) and \
               (self.type == other.type)

    def __ne__(self, other):
        return not self.__eq__(other)

    def __hash__(self):
        return (self.chrom.__hash__() ^ \
                self.chromStart ^ self.chromEnd ^ \
                self.name.__hash__() ^ \
                self.type.__hash__())

    def __repr__(self):
        return "Annotation('%s', %s, %s, '%s', '%s')" % (self.chrom, self.chromStart, self.chromEnd, self.name, self.type)

    def isOverlapping(self, other):
        if self.chrom != other.chrom:
            return False
        if is_overlapping((self.chromStart, self.chromEnd),
                          (other.chromStart, other.chromEnd)):
            return True
    

class Overlap(Base):
    '''
    table that associates overlapping annotations with each other with edges
    
    src_id, tgt_id: this overlap connects annotation id src to annotation id tgt
    '''
    __tablename__ = 'overlaps'
    
    # table fields
    src_id = Column(Integer, ForeignKey('annotations.id'), primary_key=True)
    tgt_id = Column(Integer, ForeignKey('annotations.id'), primary_key=True)
    src = relation(Annotation, primaryjoin=src_id == Annotation.id, backref='src_overlaps', collection_class=set)
    tgt = relation(Annotation, primaryjoin=tgt_id == Annotation.id, backref='tgt_overlaps', collection_class=set)
    
    def __init__(self, id1, id2):
        self.src_id = id1
        self.tgt_id = id2
    
    def __repr__(self):
        return "<Overlap('%s %s')>" % (self.src_id, self.tgt_id)    


class AnnotationDB(object):
    '''
    manages a database that stores tables of genome annotations
    '''
    def __init__(self, dburl=None, echo=True):
        '''
        constructor
    
        dburl: url or path to database.  if not specified an sqlite in-memory database is created
        '''
        # connect to database via an engine
        if (not dburl):
            dburl = 'sqlite:///:memory:'
        self.dburl = dburl
        self.engine = create_engine(dburl, echo=echo)
        
        # associate metadata with the engine
        self.metadata = Base.metadata
        self.metadata.bind = self.engine
        self.metadata.create_all()
        # self.metadata.reflect()
        # setup a session
        # TODO: explore thread local or scoped session
        self.Session = sessionmaker(bind=self.engine)
    
    def add(self, a):
        session = self.Session()
        session.add(a)
        session.commit()
        session.close()
        
    def addAll(self, annotations):
        session = self.Session()
        session.add_all(annotations)
        session.commit()
        session.close()

    def buildOverlaps(self):
        session = self.Session()
        a = aliased(Annotation)
        b = aliased(Annotation)
        q = session.query(a.id, b.id)
        q = q.filter(a.chrom == b.chrom)
        q = q.filter(not_(or_(a.chromEnd < b.chromStart,
                              a.chromStart > b.chromEnd)))
        q = q.filter(a.id < b.id)
        for overlap in q:
            session.add(Overlap(overlap[0], overlap[1]))
            session.add(Overlap(overlap[1], overlap[0]))
        session.commit()
        session.close()
    
    def getAnnotations(self, chrom=None, name=None):
        session = self.Session()
        q = session.query(Annotation)
        if name:
            q = q.filter(Annotation.name == name)
        if chrom:
            q = q.filter(Annotation.chrom == chrom)
        for annot in q:
            yield annot
        session.close()
        
    def getOverlappingPairs(self, name1, name2):
        session = self.Session()
        q = session.query(Overlap)
        q = q.filter(Overlap.src.has(name=name1))
        q = q.filter(Overlap.tgt.has(name=name2))
        for overlap in q:
            yield overlap.src, overlap.tgt
        session.close()
        return
        # TODO: this might work better (all in SQL no yielding stuff)
        q = s.query(x, y).join((e1, x.id == e1.src_id)).join((y, and_(e1.src_id == x.id, e1.tgt_id == y.id))) 

    def getCliques(self, n, names=None):
        '''
        find groups of n or more overlapping annotations
        '''        
        s = self.Session()
        # create aliases to 'n' nodes
        nodes = []
        edges = {}
        for i in xrange(n):
            nodes.append(aliased(Annotation))
        # create query that should yield sets of 'n' nodes
        q = s.query(*nodes)
        # loop for nodes 0..(n-1)
        for i, node1 in enumerate(nodes[:-1]):
            # for each index i loop from (i+1) to n
            for j, node2 in enumerate(nodes[(i+1):]):
                # join edges that originate at this node
                e = aliased(Overlap)
                q = q.join((e, and_(e.src_id == node1.id, e.tgt_id == node2.id)))
                edges[(node1, node2)] = e
                # if node1 index is zero then we still need to do the initial join to the remaining nodes
                if (i == 0):
                    q = q.join((node2, e.tgt_id == node2.id))
        
        # filter to match associated names
        if not names:
            names = []
        # extend names to match 'n' (size of cliques formed)
        for i in xrange(len(names), n):
            names.append('*')
        # filter to constrain node names
        free_cols = []
        for i, (name, node) in enumerate(zip(names, nodes)):
            if name == '*':
                free_cols.append(node)
                continue
            q = q.filter(node.name == name)
        # if unrestricted (free) columns remain, impose an arbitrary order to remove non-unique combinations
        for node1, node2 in zip(free_cols, free_cols[1:]):
            q = q.filter(node1.id < node2.id)

        # return the annotations
        for annot in q:
            yield annot

        s.close()
        #return
        #return q
    
        # TODO: it appears that using these complete node constraints does not speed up the query, so for now the
        # function returns above
        # loop for nodes 0..n to build node constraints
#        for i, node1 in enumerate(nodes):
#            # the first node is automatically joined as part of the query so skip it
#            if i == 0:
#                continue            
#            join_conds = []
#            # node1 is source
#            for node2 in nodes[(i+1):]:
#                e = edges[(node1, node2)]
#                join_conds.append(e.src_id == node1.id)
#                join_conds.append(node1.id < node2.id)
#            # node1 is target
#            for node2 in nodes[:i]:
#                e = edges[(node2, node1)]
#                join_conds.append(e.tgt_id == node1.id)
#            # AND together all the join conditions
#            join_conds = and_(*join_conds)
#            # perform the join
#            q = q.join((node1, join_conds))
        # example query for 4
        #q3 = s.query(w, x, y, z).join((x, and_(w.id < x.id, e1.tgt_id == x.id))).
        #join((y, and_(x.id < y.id, e2.tgt_id == y.id))).join((z, and_(y.id < z.id, e3.tgt_id == z.id))).
        #join((e1, w.id == e1.src_id)).join((e2, w.id == e2.src_id)).join((e3, w.id == e3.src_id)).
        #join((e4, and_(x.id == e4.src_id, y.id == e4.tgt_id))).join((e5, and_(x.id == e5.src_id, z.id == e5.tgt_id))).
        #join((e6, and_(y.id == e6.src_id, z.id == e6.tgt_id)))

    def getOverlapHardCoded(self, names):
        s = self.Session()

        a = aliased(Overlap)
        b = aliased(Overlap)
        c = aliased(Overlap)
        d = aliased(Overlap)
        
        q_erg = s.query(Annotation.id, b.tgt_id, c.tgt_id, d.tgt_id).filter(Annotation.name=='ERG_peak')
        q_erg = q_erg.join((b, Annotation.src_overlaps), (c, Annotation.src_overlaps), (d, Annotation.src_overlaps))
        q_erg = q_erg.filter(b.tgt.has(name='DNApkcs_peak')).filter(c.tgt.has(name='Ku80_peak')).filter(d.tgt.has(name='pH2AX_peak'))

        #q_erg = s.query(Annotation.id, b.tgt_id, c.tgt_id, d.tgt_id).filter(Annotation.name=='ERG_peak').join((b, Annotation.src_overlaps), (c, Annotation.src_overlaps), (d, Annotation.src_overlaps)).filter(b.tgt.has(name='DNApkcs_peak')).filter(c.tgt.has(name='Ku80_peak')).filter(d.tgt.has(name='pH2AX_peak'))

        q_dna = s.query(a.tgt_id, Annotation.id, c.tgt_id, d.tgt_id).filter(Annotation.name=='DNApkcs_peak')
        q_dna = q_dna.join((a, Annotation.src_overlaps), (c, Annotation.src_overlaps), (d, Annotation.src_overlaps))
        q_dna = q_dna.filter(a.tgt.has(name='ERG_peak')).filter(c.tgt.has(name='Ku80_peak')).filter(d.tgt.has(name='pH2AX_peak'))

        q_ku80 = s.query(a.tgt_id, b.tgt_id, Annotation.id, d.tgt_id).filter(Annotation.name=='Ku80_peak')
        q_ku80 = q_ku80.join((a, Annotation.src_overlaps), (b, Annotation.src_overlaps), (d, Annotation.src_overlaps))
        q_ku80 = q_ku80.filter(a.tgt.has(name='ERG_peak')).filter(b.tgt.has(name='DNApkcs_peak')).filter(d.tgt.has(name='pH2AX_peak'))

        q_ph2ax = s.query(a.tgt_id, b.tgt_id, c.tgt_id, Annotation.id).filter(Annotation.name=='pH2AX_peak')
        q_ph2ax = q_ph2ax.join((a, Annotation.src_overlaps), (b, Annotation.src_overlaps), (c, Annotation.src_overlaps))
        q_ph2ax = q_ph2ax.filter(a.tgt.has(name='ERG_peak')).filter(b.tgt.has(name='DNApkcs_peak')).filter(c.tgt.has(name='Ku80_peak'))

        erg = set(q_erg.all())
        dna = set(q_dna.all())
        ku80 = set(q_ku80.all())
        ph2ax = set(q_ph2ax.all())

        print 'erg:', len(erg)
        print 'dna:', len(dna)
        print 'ku80:', len(ku80)
        print 'ph2ax:', len(ph2ax)

        combined = set([])
        combined.update(erg)
        combined.intersection_update(dna)
        print 'erg intersect dna:', len(combined)
        combined.intersection_update(ku80)
        print 'erg intersect dna intersect ku80:', len(combined)
        combined.intersection_update(ph2ax)
        print 'erg intersect dna intersect ku80 intersect ph2ax:', len(combined)
        
        x = set([])
        x.update(erg)
        x.intersection_update(ph2ax)
        print 'erg intersect ph2ax:', len(x)
        
        return combined


    def getJumboo(self, names):
        session = self.Session()
        # find records where (len(names) - 1) or more annotations overlap
        q = session.query(Overlap.src_id)
        q = q.group_by(Overlap).having(func.count(Overlap.src_id) >= (len(names) - 1))
        len_stmt = q.subquery()
        # convert the areas of high overlap to their matching annotations
        q = session.query(Annotation).join((len_stmt, Annotation.id == len_stmt.c.src_id))
        stmt = q.subquery()
        annot_alias = aliased(Annotation, stmt)
        # query against all requested names
        q1 = session.query(annot_alias).filter_by(name=names[0])
        q2 = q1.filter(Annotation.tgt_overlaps.any(Overlap.src.has(name=names[1])))
        
        # TODO: right now I am just selecting annotations with name 'ERG' that happen to overlap
        # with annotations 'DNApkcs' and 'Ku80', etc.
        z = s.query(Annotation).select_from(stmt)
        
        q1 = s.query(x, y, z).filter(x.tgt_id == y.src_id).filter(y.tgt_id == z.src_id).filter(z.tgt_id == x.src_id)

        w = s.query(Overlap).filter(Overlap.src.has(name='ERG_peak'))

        w = s.query(Annotation).filter(Annotation.name == 'ERG_peak').filter(Annotation.src_overlaps.any())
        x = s.query(Annotation).filter(Annotation.name == 'DNApkcs_peak').filter(Annotation.src_overlaps.any())
        y = same
        z = same
        
        s.query(Annotation).filter(Annotation.tgt_overlaps.any(Overlap.src.has()))

        q = s.query(Annotation).filter(Annotation.name=='a').filter(Annotation.tgt_overlaps.any(Overlap.src.has(name='b'))).all()
        q1 = q.filter(Annotation.tgt_overlaps.any(Overlap.src.has(name='DNApkcs_peak')))                                                   

        z = s.query(Annotation).join((Overlap, Annotation.src_overlaps)).group_by(Overlap.src_id).having(func.count(Overlap.src_id) >= 3)
        m = s.query(Annotation, Overlap).join((Overlap, Annotation.src_overlaps)).filter(Overlap.tgt.has(name='DNApkcs_peak'))

        # amazing command and super fast
        m = s.query(Annotation, Annotation.id, a.tgt_id, b.tgt_id).join((a, Annotation.src_overlaps), (b, Annotation.src_overlaps)).filter(a.tgt.has(name='DNApkcs_peak')).filter(b.tgt.has(name='Ku80_peak'))

        m = s.query(Annotation, Annotation.id, a.tgt_id, b.tgt_id, c.tgt_id).join((a, Annotation.src_overlaps), (b, Annotation.src_overlaps), (c, Annotation.src_overlaps)).filter(a.tgt.has(name='DNApkcs_peak')).filter(b.tgt.has(name='Ku80_peak')).filter(c.tgt.has(name='pH2AX_peak'))

        # garbage
        print(s.query(a1, a2).join((a2, a1.src_overlaps.any() == a2.tgt_overlaps.any())))
        s.query(a1, a2, o1).filter(a1.name=='a').filter(a2.name=='b').join((o1, and_(a1.id == o1.src_id, a2.id == o1.tgt_id))).all()

        # works!!
        q_erg = s.query(Annotation.id, annot2, c.tgt_id, d.tgt_id).filter(Annotation.name=='ERG_peak').join((b, Annotation.src_overlaps), (c, Annotation.src_overlaps), (d, Annotation.src_overlaps)).filter(b.tgt.has(name='DNApkcs_peak')).filter(c.tgt.has(name='Ku80_peak')).filter(d.tgt.has(name='pH2AX_peak')).join((annot2, annot2.id == b.tgt_id))

        # works!!!
        q_erg = s.query(Annotation, annot2, annot3, annot4).filter(Annotation.name=='ERG_peak').join((b, Annotation.src_overlaps), (c, Annotation.src_overlaps), (d, Annotation.src_overlaps)).filter(b.tgt.has(name='DNApkcs_peak')).filter(c.tgt.has(name='Ku80_peak')).filter(d.tgt.has(name='pH2AX_peak')).join((annot2, annot2.id == b.tgt_id), (annot3, annot3.id == c.tgt_id), (annot4, annot4.id == d.tgt_id))

        # works!!! without annotations at all.. just a query of edges
        # finds all unique "triangles"
        o1 = aliased(Overlap)
        o2 = aliased(Overlap)
        o3 = aliased(Overlap)
        s.query(o1, o2, o3).join((o2, and_(o1.tgt_id == o2.src_id, o1.src_id < o2.src_id)), (o3, and_(o2.tgt_id == o3.src_id, o2.src_id < o3.src_id, o3.tgt_id == o1.src_id))).all()


    def test(self):
        s = self.Session()
        s.add(Annotation('chr1', 0, 100, 'a', 'a'))
        s.add(Annotation('chr1', 50, 55, 'b', 'b'))
        s.add(Annotation('chr1', 80, 120, 'b', 'b'))
        s.add(Annotation('chr1', 150, 200, 'a', 'a'))
        s.add(Annotation('chr1', 140, 160, 'b', 'b'))
        s.add(Annotation('chr1', 300, 500, 'a', 'a'))
        s.add(Annotation('chr1', 40, 60, 'c', 'c'))
        s.add(Annotation('chr1', 400, 600, 'c', 'c'))
        s.add(Annotation('chr1', 75, 85, 'c', 'c'))
        s.add(Annotation('chr1', 1000, 2000, 'c', 'c'))
        s.commit()
        s.close()
    
    def test2(self):
        s = self.Session()
        s.add(Annotation('chr1', 0, 100, 'a', 'a'))
        s.add(Annotation('chr1', 1000, 2000, 'a', 'a'))
        s.add(Annotation('chr1', 3000, 4000, 'a', 'a'))
        s.add(Annotation('chr1', 50, 150, 'b', 'b'))
        s.add(Annotation('chr1', 170, 180, 'b', 'b'))        
        s.add(Annotation('chr1', 10000, 20000, 'b', 'b'))
        s.add(Annotation('chr1', 30000, 40000, 'b', 'b'))
        s.add(Annotation('chr1', 75, 125, 'c', 'c'))
        s.add(Annotation('chr1', 150, 250, 'c', 'c'))
        s.add(Annotation('chr1', 5000, 15000, 'c', 'c'))
        s.add(Annotation('chr1', 80, 85, 'd', 'd'))
        s.add(Annotation('chr1', 10, 20, 'd', 'd'))
        s.commit()
        s.close()        

    

    



